import asyncio
import logging
from typing import Optional
from contextlib import asynccontextmanager
import motor.motor_asyncio
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
from config.settings import settings

logger = logging.getLogger(__name__)

class MongoConnectionManager:
    """MongoDB连接管理器，处理连接池、重连和健康检查"""
    
    def __init__(self, host: str = "localhost", port: int = 27017, database_name: str = "site_evaluation"):
        self.host = host
        self.port = port
        self.database_name = database_name
        self.client: Optional[motor.motor_asyncio.AsyncIOMotorClient] = None
        self.database = None
        self.is_connected = False
        self.connection_url = f"mongodb://{host}:{port}"
        self.max_retries = 3
        self.retry_delay = 2.0
    
    async def connect(self) -> bool:
        """建立数据库连接"""
        if self.is_connected and self.client:
            return True
            
        try:
            logger.info(f"正在连接MongoDB: {self.connection_url}")
            
            # 配置连接选项
            client_options = {
                'serverSelectionTimeoutMS': 5000,  # 5秒超时
                'connectTimeoutMS': 10000,         # 10秒连接超时
                'maxPoolSize': 10,                 # 最大连接池大小
                'minPoolSize': 1,                  # 最小连接池大小
                'maxIdleTimeMS': 30000,            # 连接空闲时间
                'retryWrites': True,               # 启用写重试
            }
            
            self.client = motor.motor_asyncio.AsyncIOMotorClient(
                self.connection_url, 
                **client_options
            )
            
            # 测试连接
            await self._test_connection()
            
            self.database = self.client[self.database_name]
            self.is_connected = True
            
            logger.info(f"MongoDB连接成功: {self.database_name}")
            return True
            
        except Exception as e:
            logger.error(f"MongoDB连接失败: {e}")
            await self._cleanup_connection()
            return False
    
    async def _test_connection(self) -> None:
        """测试数据库连接"""
        if not self.client:
            raise ConnectionFailure("客户端未初始化")
            
        try:
            # 执行ping命令测试连接
            await self.client.admin.command('ping')
            
            # 获取服务器信息
            server_info = await self.client.server_info()
            logger.info(f"MongoDB服务器版本: {server_info.get('version', 'unknown')}")
            
        except (ConnectionFailure, ServerSelectionTimeoutError) as e:
            raise ConnectionFailure(f"数据库连接测试失败: {e}")
    
    async def disconnect(self) -> None:
        """断开数据库连接"""
        if not self.client:
            return
            
        try:
            logger.info("正在断开MongoDB连接")
            await self._cleanup_connection()
            logger.info("MongoDB连接已断开")
        except Exception as e:
            logger.error(f"断开连接时出错: {e}")
    
    async def _cleanup_connection(self) -> None:
        """清理连接资源"""
        if self.client:
            self.client.close()
            self.client = None
        self.database = None
        self.is_connected = False
    
    async def ensure_connection(self) -> bool:
        """确保数据库连接可用，如果断开则重连"""
        if not self.is_connected or not self.client:
            return await self.connect()
        
        try:
            # 快速健康检查
            await self.client.admin.command('ping')
            return True
        except Exception as e:
            logger.warning(f"连接健康检查失败，尝试重连: {e}")
            await self._cleanup_connection()
            return await self.connect()
    
    async def retry_connection(self) -> bool:
        """重试连接，带指数退避"""
        for attempt in range(self.max_retries):
            if attempt > 0:
                delay = self.retry_delay * (2 ** (attempt - 1))
                logger.info(f"第{attempt + 1}次重连尝试，等待{delay}秒...")
                await asyncio.sleep(delay)
            
            if await self.connect():
                return True
                
        logger.error(f"经过{self.max_retries}次尝试后连接失败")
        return False
    
    @asynccontextmanager
    async def get_database(self):
        """获取数据库实例的上下文管理器"""
        if not await self.ensure_connection():
            raise ConnectionFailure("无法建立数据库连接")
        
        try:
            yield self.database
        except Exception as e:
            logger.error(f"数据库操作出错: {e}")
            raise
    
    @asynccontextmanager
    async def get_collection(self, collection_name: str):
        """获取集合实例的上下文管理器"""
        async with self.get_database() as db:
            yield db[collection_name]
    
    async def create_indexes(self, collection_name: str = "articles") -> bool:
        """创建数据库索引"""
        if not await self.ensure_connection():
            logger.error("无法连接数据库，跳过索引创建")
            return False
        
        try:
            async with self.get_collection(collection_name) as collection:
                # 为URL创建唯一索引，避免重复抓取
                await collection.create_index("url", unique=True)
                logger.info("创建URL唯一索引成功")
                
                # 为行业关键词和爬取时间创建复合索引
                await collection.create_index([
                    ("industry_keyword", 1),
                    ("crawl_time", -1)
                ])
                logger.info("创建行业关键词+爬取时间复合索引成功")
                
                # 为域名创建索引
                await collection.create_index("source_domain")
                logger.info("创建来源域名索引成功")
                
                # 为状态创建索引
                await collection.create_index("status")
                logger.info("创建状态索引成功")
                
                logger.info("所有数据库索引创建完成")
                return True
                
        except Exception as e:
            logger.error(f"创建索引失败: {e}")
            return False
    
    async def get_connection_status(self) -> dict:
        """获取连接状态信息"""
        status = {
            "connected": self.is_connected,
            "host": self.host,
            "port": self.port,
            "database": self.database_name,
            "connection_url": self.connection_url
        }
        
        if self.is_connected and self.client:
            try:
                server_info = await self.client.server_info()
                status.update({
                    "server_version": server_info.get('version'),
                    "uptime": server_info.get('uptime'),
                })
            except Exception as e:
                status["error"] = str(e)
        
        return status

# 全局连接管理器实例
connection_manager = MongoConnectionManager(
    host="localhost",
    port=27017,
    database_name=settings.DATABASE_NAME
) 